#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Author  :   FRzhuizhu
@License :   (C) Copyright 2013-2017, FRzhuizhu
@Contact :   736019332@qq.com
@Software:   PyCharm
@File    :   rebuilt.py
@Time    :   2020/9/11 20:22
@Desc    :
'''

import json
import os
import logging
import logging.config
import sys
import shutil
import time
import math


class Vec:

    def __init__(self, data: list):
        self.x, self.y, self.z = data

    def dot(self, v):
        return self.x * v.x + self.y * v.y + self.z * v.z

    def angle(self, v):
        return math.acos(self.dot(v) / (self.mol() * v.mol())) * 180 / math.pi

    def mol(self):
        return (self.x ** 2 + self.y ** 2 + self.z ** 2) ** 0.5

    def data(self):
        return [self.x, self.y, self.z]

    def __str__(self):
        return str(self.data())

    def __repr__(self):
        return f'<Vec:{self.data()}>'


class Params:

    def __init__(self, data):
        # 为了编写代码时可以自动补全参数名
        self.ms_install_path = None
        self.np = None
        self.BASFILE_path = None
        self.start_path = None
        self.input_path = None
        self.output_path = None

        self.__dict__.update(data.copy())


class Until:

    def _transformCore(self, strout_lines):
        v_a, v_b, v_c = Vec([eval(num) for num in strout_lines[0].split()]), \
                        Vec([eval(num) for num in strout_lines[1].split()]), \
                        Vec([eval(num) for num in strout_lines[2].split()])
        atom_dict = {}
        for line in strout_lines[6:]:
            # 2020-9-17 21:10:48 添加忽略Vac行的功能
            if 'Va' in line:
                continue
            x, y, z, atom = line.split()
            if atom_dict.get(atom, None):
                atom_dict[atom].append([eval(num) for num in (x, y, z)])
            else:
                atom_dict[atom] = [[eval(num) for num in (x, y, z)]]
        DATE = time.asctime()
        car_lines = [
            '!BIOSYM archive 3',
            'PBC=ON',
            '',
            f'!DATE {DATE}',
            'PBC' + ('{: >10.4f}' * 6).format(v_a.mol(), v_b.mol(), v_c.mol(), v_a.angle(v_b), v_a.angle(v_c),
                                              v_b.angle(v_c)),
        ]
        for atom, points in atom_dict.items():
            for index, point in enumerate(points):
                car_lines.append(f'{atom + str(index + 1): <5}' +
                                 ('{: >15.9f}' * 3).format(*point) +
                                 ' XXXX 1      xx      'f'{atom: <4}0.000')
        car_lines.append('end')
        car_lines.append('end\n')
        return car_lines

    def transform2car(self, strout_path, output_path=None, retext=False):
        # 提供 str.out 转化为 *.car 的外部调用接口
        with open(strout_path, 'r', encoding='utf-8') as f:
            strout_lines = f.readlines()
        out_car_text = '\n'.join(self._transformCore(strout_lines))
        if output_path:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(out_car_text)
        if retext:
            return out_car_text


class AutoMS:

    def __init__(self, enable_conf=True):
        self.logger = self._init_logger()
        self.enable_conf = enable_conf
        self.input_params = self._init_params()
        self._add_env()
        self.finish = self._init_finished_list()
        self.until = Until()
        self.runing_path = os.path.abspath('.')

    def _init_logger(self):
        logging.config.fileConfig('logger.conf')
        logger = logging.getLogger('logger')
        logger.info('{:=^50}'.format('脚本已启动，日志模块初始化完毕'))
        return logger

    def _init_params(self):
        params = {
            "ms_install_path": "不启用",
            "np": os.cpu_count(),
            "BASFILE_path": os.path.abspath(os.path.join('.', 'BASFILE_v3.5')),
            "start_path": os.path.abspath('.'),
            "input_path": os.path.abspath([file for file in os.listdir('.') if os.path.splitext(file)[-1]=='.input'][0]),
            "output_path": os.path.abspath(os.path.join('.', 'output'))
        }
        if self.enable_conf:
            if not os.path.exists('conf.json'):
                with open('conf.json', 'w', encoding='utf-8') as f:
                    json.dump(params, fp=f, ensure_ascii=False, indent=True)
            else:
                with open('conf.json', 'r', encoding='utf-8') as f:
                    params = json.load(f)
        assert params['ms_install_path'], 'MS目录不能为空'
        self.logger.info('[使用参数]MS目录：%s', params['ms_install_path'])

        assert params['BASFILE_path'], ' '
        self.logger.info('[使用参数]BASFILE目录：%s', params['BASFILE_path'])

        assert params['start_path'], 'start目录不能为空'
        self.logger.info('[使用参数]start目录：%s', params['start_path'])

        assert params['input_path'], 'input目录不能为空'
        self.logger.info('[使用参数]input目录：%s', params['input_path'])

        assert params['output_path'], 'output目录不能为空'
        self.logger.info('[使用参数]output目录：%s', params['output_path'])

        if not params.get('np', 0) or params['np'] < 0:
            params['np'] = os.cpu_count()
        # params['np'] = params.get('np', os.cpu_count())
        self.logger.info('[使用参数]启用核心数：%d', params['np'])
        return Params(params)

    def _add_env(self):
        if self.enable_conf:
            runbat_path = os.path.join(self.input_params.ms_install_path, 'etc', 'DMol3', 'bin')
            self.logger.info('添加临时环境变量：%s', runbat_path)
            sys.path.append(runbat_path)
        else:
            self.logger.info('不设置环境变量')

    def _init_finished_list(self):
        finished = {
            self.input_params.start_path: [],
        }
        if os.path.exists('finishedList.json'):
            with open('finishedList.json', 'r', encoding='utf-8') as f:
                finished = json.load(f)
        else:
            with open('finishedList.json', 'w', encoding='utf-8') as f:
                json.dump(finished, f, ensure_ascii=False, indent=True)
        return finished

    def _sava_finished_list(self):
        os.chdir(self.runing_path)
        with open('finishedList.json', 'w', encoding='utf-8') as f:
            json.dump(self.finish, f, ensure_ascii=False, indent=True)

    def _save_error(self, file):
        os.chdir(self.runing_path)
        with open('error.txt', 'a+', encoding='utf-8') as f:
            f.write(file + '\n')

    def _copy_input_file(self, seedName):
        os.chdir(self.input_params.output_path)
        if os.path.exists(seedName): shutil.rmtree(seedName)
        os.mkdir(seedName)
        self.logger.info('进入目录:%s', os.path.abspath(seedName))
        os.chdir(seedName)
        with open(os.path.join(self.input_params.start_path, seedName, 'str.out'), 'r', encoding='utf-8') as f:
            strout_lines = f.readlines()
        with open(f'task_{seedName}.car', 'w', encoding='utf-8') as f:
            f.write('\n'.join(self.until._transformCore(strout_lines)))
        self.logger.info('生成car文件:%s', os.path.abspath(f'task_{seedName}.car'))
        res_path = shutil.copy(self.input_params.input_path, f'task_{seedName}.input')
        self.logger.info('复制文件:%s', os.path.abspath(res_path))
        res_path = shutil.copy(self.input_params.BASFILE_path, '.')
        self.logger.info('复制文件:%s', os.path.abspath(res_path))

    def _get_result(self, seedName):
        # 失败返回False， 成功返回Total Energy: str
        os.chdir(os.path.join(self.input_params.output_path, seedName))
        if not os.path.exists(f'task_{seedName}.outmol'):
            self.logger.error('outmol文件不存在:%s', os.path.abspath(f'task_{seedName}.outmol'))
            return False
        with open(f'task_{seedName}.outmol', 'r', encoding='utf-8') as f:
            outmol_lines_R = f.readlines()[::-1]
        if any(['DMol3 job finished successfully' in l for l in outmol_lines_R]):
            for l in outmol_lines_R:
                if l.split()[0] == 'Ef':
                    return l.split()[1][:-2]
                elif l.split()[0] == 'opt==':
                    return l.split()[2]
            else:
                self.logger.error('outmol中找不到Total Energy')
                return False
        else:
            return False

    def _sava_totalenergy(self, seedName, value):
        os.chdir(os.path.join(self.input_params.start_path, seedName))
        with open('energy', 'w', encoding='utf-8') as f:
            f.write(value)

    def _runbat(self, seedName):
        os.chdir(os.path.join(self.input_params.output_path, seedName))
        cmd = f'RunDMol3.{"sh" if "linux" in sys.platform else "bat"} -np {self.input_params.np} task_{seedName}'
        self.logger.info('开启任务:%s', cmd)
        self.logger.info('runDmol3返回:%s', os.system(cmd))

    def _is_ready(self, seedName):
        os.chdir(os.path.join(self.input_params.start_path, seedName))
        return os.path.exists('str.out')

    def start_listen(self):
        try:
            while True:
                os.chdir(self.input_params.start_path)
                os.system('touch ready')
                time.sleep(3)
                update_dirs = [d for d in os.listdir(self.input_params.start_path)
                               if os.path.isdir(d) and d.isdigit()
                               and d not in self.finish.get(self.input_params.start_path, [])]

                for d in update_dirs:
                    if not self._is_ready(seedName=d):
                        continue
                    self._copy_input_file(seedName=d)
                    self._runbat(seedName=d)
                    res = self._get_result(seedName=d)
                    if res:
                        self._sava_totalenergy(seedName=d, value=res)
                    else:
                        self.logger.error('优化失败！')
                        self._save_error(os.path.join(self.input_params.output_path, d))
                    self.finish.get(self.input_params.start_path).append(d)
                    self._sava_finished_list()


        except Exception as e:
            self.logger.exception(e)
            raise e


if __name__ == '__main__':
    app = AutoMS(enable_conf=False)
    app.start_listen()

    # str.out 转化 *.car 的调用接口
    # Until().transform2car(strout_path='', output_path='', retext=False) -> out_car_text: str
    # Until().transform2car('str.out', output_path='1.car')
